{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "cEWCN4AIkg4V",
        "SPbJdMWrkoew",
        "nHhkCMPKk3u0",
        "v-mPIBULlDLj",
        "oEFSKx5BlVRb",
        "2o-0zuSCliJp",
        "wJEyV68IlpGs",
        "tBCo8vtol-gT",
        "4fYYD5p_njNM",
        "7p6vlMOaoBjK",
        "577k5MrvoEz6"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Ankur-Samanta/Multimodal_Transformer_CardiacAbnormailityECGClassification/blob/main/MVMTnet_Model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "MVMTnet: Multi-variate Multi-modal Transformer Network for Multi-class Classification of Cardiac Irregularities Using ECG Waveforms and Clinical Notes\n",
        "\n",
        "Authors: Ankur Samanta, Meghna Ravikumar, Mark Karlov, Christian Clarke"
      ],
      "metadata": {
        "id": "90vauTsN_PKb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Notebook Initialization"
      ],
      "metadata": {
        "id": "CBkzUw76dmnB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Connect to Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive', force_remount=True)"
      ],
      "metadata": {
        "id": "SVskVhGqeBdP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Required Dependencies and Imports"
      ],
      "metadata": {
        "id": "MHx7lO92eSMI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# ECG Dependencies\n",
        "!pip install ecg_plot\n",
        "!pip install pc\n",
        "!pip install wfdb\n",
        "\n",
        "# Bert Dependencies\n",
        "!pip install transformers"
      ],
      "metadata": {
        "id": "MnlQARqPeDro"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Required Imports\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import csv\n",
        "\n",
        "import math\n",
        "import random\n",
        "from typing import Tuple\n",
        "import copy\n",
        "import ast\n",
        "\n",
        "import pc\n",
        "import os\n",
        "\n",
        "import ecg_plot\n",
        "import wfdb\n",
        "from wfdb.io.record import Record, rdrecord\n",
        "from wfdb.plot.plot import plot_wfdb\n",
        "\n",
        "from socket import socket\n",
        "\n",
        "import torch\n",
        "from torch import nn, Tensor\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.utils.data.sampler import SubsetRandomSampler\n",
        "from torch.nn import TransformerEncoder, TransformerEncoderLayer\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "import torchvision\n",
        "from torchvision import datasets,transforms\n",
        "\n",
        "from transformers import BertTokenizer, BertModel"
      ],
      "metadata": {
        "id": "uZ-rISe4ed6H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Get the Data"
      ],
      "metadata": {
        "id": "jesxH_DMeIfo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Unzip the PTB-XL dataset\n",
        "\n",
        "# Ankur's Path:\n",
        "!unzip \"/content/gdrive/MyDrive/Colab Notebooks/CardiacAbnormalityTransformerProject/physionet.org/ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.2.zip\"\n",
        "# Meghna's Path:\n",
        "!unzip \"/content/gdrive/MyDrive/Fall 2022/APS360/CardiacAbnormalityTransformerProject/physionet.org/ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.2.zip\"\n",
        "# Mark's path\n",
        "!unzip \"/content/gdrive/MyDrive/CardiacAbnormalityTransformerProject/physionet.org/ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.2.zip\""
      ],
      "metadata": {
        "id": "P5UqVyFEeQJ_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Unzip the clinical embeddings\n",
        "\n",
        "# Chris\n",
        "!unzip \"/content/gdrive/MyDrive/CardiacAbnormalityTransformerProject/clinicalEmbeddings.zip\"\n",
        "\n",
        "# Ankur\n",
        "!unzip \"/content/gdrive/MyDrive/Colab Notebooks/CardiacAbnormalityTransformerProject/clinicalEmbeddings.zip\"\n",
        "\n",
        "#Meghna\n",
        "!unzip \"/content/gdrive/MyDrive/Fall 2022/APS360/CardiacAbnormalityTransformerProject/clinicalEmbeddings.zip\""
      ],
      "metadata": {
        "id": "Ncopz6VffiGz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Data Preprocessing"
      ],
      "metadata": {
        "id": "cuVtz-iPfy90"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def load_raw_data(df, sampling_rate, path):\n",
        "    \"\"\"\n",
        "    The first line describes concisely what the functions does. \n",
        "    \n",
        "    Args: \n",
        "        argument1 (str): Description of argument1. If you have lots to say about\n",
        "            a function you can break to a new line and indent the text so it fits. \n",
        "        argument2 (int, optional): Description of argument2. \n",
        "    \n",
        "    Returns: \n",
        "        str: Optional description explaining the value returned by the function. \n",
        "        \n",
        "    Raises:\n",
        "        ValueError: Description of any errors or exceptions intentionally raised. \n",
        "    \n",
        "    Notes: \n",
        "        Credits: https://physionet.org/content/ptb-xl/1.0.1/example_physionet.py\n",
        "    \"\"\"\n",
        "    if sampling_rate == 100:\n",
        "        # Loading all data with signal and meta information\n",
        "        data = [wfdb.rdsamp(path+f)[0] for f in df.filename_lr]\n",
        "    else:\n",
        "        data = [wfdb.rdsamp(path+f) for f in df.filename_hr]\n",
        "    data = np.array(data)\n",
        "    return data"
      ],
      "metadata": {
        "id": "aErl-koffzZr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load raw signal data\n",
        "\n",
        "# Ankur/Mark/Chris's Path\n",
        "df = pd.read_csv('/content/gdrive/MyDrive/Colab Notebooks/CardiacAbnormalityTransformerProject/undersampledY')\n",
        "\n",
        "# Meghna's path\n",
        "df = pd.read_csv('/content/gdrive/MyDrive/Fall 2022/APS360/CardiacAbnormalityTransformerProject/undersampledY')\n",
        "\n",
        "X = load_raw_data(df, 100, '/content/ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.2/')"
      ],
      "metadata": {
        "id": "0QUhqfXh-Pe4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "10d73501-ca94-4364-88f3-2869e542d022"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(11271, 1000, 12)"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Sanity check: see how many ecg segments have been loaded\n",
        "X.shape"
      ],
      "metadata": {
        "id": "vueZflJ6habn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# note: rename this class to something better \n",
        "class M(Dataset):\n",
        "    def __init__(self, diagnostics_superclass, X):\n",
        "        \"\"\"\n",
        "        Docstring goes here\n",
        "        \"\"\"\n",
        "        self.labels = diagnostics_superclass\n",
        "        self.data = X\n",
        "        self.id = df.index.to_list()\n",
        "        self.notes = df.report.to_list()\n",
        "        self.map = {'[\\'NORM\\']':0,'[\\'MI\\']':1,'[\\'STTC\\']':2,'[\\'CD\\']':3,'[\\'HYP\\']':4, '[\\'HYP\\', \\'MI\\']':5, '[\\'HYP\\', \\'CD\\']':6,\n",
        "                    '[\\'HYP\\', \\'STTC\\']':7,'[\\'MI\\', \\'CD\\']':8,'[\\'MI\\', \\'STTC\\']':9,'[\\'CD\\', \\'STTC\\']':10,'[\\'HYP\\', \\'MI\\', \\'STTC\\']':11,\n",
        "                    '[\\'HYP\\', \\'MI\\', \\'CD\\']':12, '[\\'MI\\', \\'CD\\', \\'STTC\\']':13,'[\\'HYP\\', \\'CD\\', \\'STTC\\']':14,'[\\'HYP\\', \\'MI\\', \\'CD\\', \\'STTC\\']':15, \n",
        "                    '[\\'CD\\', \\'NORM\\']':16, '[\\'CD\\', \\'NORM\\', \\'STTC\\']':17, '[\\'NORM\\', \\'STTC\\']':18, '[\\'HYP\\', \\'NORM\\']':19, '[\\'MI\\', \\'NORM\\']':20, \n",
        "                    '[\\'HYP\\', \\'MI\\', \\'NORM\\']':21, '[\\'HYP\\', \\'CD\\', \\'NORM\\']':22, '[\\'HYP\\', \\'MI\\', \\'CD\\', \\'NORM\\']':23, '[\\'HYP\\', \\'MI\\', \\'CD\\', \\'NORM\\', \\'STTC\\']':24,\n",
        "                    '[\\'HYP\\', \\'NORM\\', \\'STTC\\']':25, '[\\'MI\\', \\'NORM\\', \\'STTC\\']':26, '[\\'HYP\\', \\'MI\\', \\'NORM\\', \\'STTC\\']':27, '[\\'HYP\\', \\'CD\\', \\'NORM\\', \\'STTC\\']':28, '[\\'MI\\', \\'CD\\', \\'NORM\\', \\'STTC\\']':29}\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        \"\"\"\n",
        "        Docstring goes here\n",
        "        \"\"\"\n",
        "        return self.data[i][:250] , self.map[\"\".join(self.labels[i])], torch.load(\"/content/\"+ str(self.id[i]) + \".tensor\") \n",
        "\n",
        "    def __len__(self):\n",
        "        \"\"\"\n",
        "        Docstring goes here\n",
        "        \"\"\"\n",
        "        return len(self.labels)"
      ],
      "metadata": {
        "id": "DqjZ_ANqhkBA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_data_loader(train_data, val_data, batch_size):\n",
        "  \"\"\"\n",
        "  Docstring goes here\n",
        "  \"\"\"\n",
        "  return torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True), torch.utils.data.DataLoader(val_data, batch_size=batch_size, shuffle=True)"
      ],
      "metadata": {
        "id": "cTG6xuyUiVb0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_split_data(dataset):\n",
        "  \"\"\"\n",
        "  Docstring goes here\n",
        "  \"\"\"\n",
        "  train_size = int(0.1 * len(dataset))\n",
        "  valid_size = int(0.1 * len(dataset))\n",
        "  test_size = len(dataset) - train_size - valid_size\n",
        "  train_loader, val_loader, test_loader = torch.utils.data.random_split(dataset, [train_size, valid_size, test_size])\n",
        "  return train_loader, val_loader, test_loader"
      ],
      "metadata": {
        "id": "BrH3Q6ujiXGL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Comment goes here\n",
        "diagnostics_superclass = df['diagnostic_superclass'].to_list()\n",
        "diagnostics_superclass = np.array(diagnostics_superclass)\n",
        "ds = M(diagnostics_superclass, X)"
      ],
      "metadata": {
        "id": "7ap-VYjRiGs0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Get and load data\n",
        "train_data, val_data, test_data = get_split_data(ds)\n",
        "train_loader, val_loader = get_data_loader(train_data, val_data, 1)"
      ],
      "metadata": {
        "id": "kOMxdLaOit4q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Sanity check: see how many samples are in the training set\n",
        "print(len(train_loader))"
      ],
      "metadata": {
        "id": "eIt4OZsGi3Bq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Model Utility Functions"
      ],
      "metadata": {
        "id": "fJWth-NQjChy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_clones(module, N):\n",
        "    \"\"\"\n",
        "    Creates clones of N encoder and decoder layers.\n",
        "    \n",
        "    Args: \n",
        "        argument1 (str): Description of argument1.\n",
        "        argument2 (int): Description of argument2. \n",
        "    \n",
        "    Returns: \n",
        "        str: Optional description explaining the value returned by the function. \n",
        "    \n",
        "    Notes: \n",
        "        Credits: \n",
        "    \"\"\"\n",
        "    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])"
      ],
      "metadata": {
        "id": "BDu4CMAzjDuw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def scaled_dot_product_attention(query, key, value, mask=None):\n",
        "    \"\"\"\n",
        "    query, key, value : batch_size * heads * max_len * d_h\n",
        "    used in multihead attention class\n",
        "    \n",
        "    Args: \n",
        "        argument1 (int): Description of argument1.\n",
        "        argument2 (int): Description of argument2. \n",
        "    \n",
        "    Returns: \n",
        "        str: batch_size * heads * max_len * d_h \n",
        "    \n",
        "    Notes: \n",
        "        Credits: \n",
        "    \"\"\"\n",
        "    matmul = torch.matmul(query,key.transpose(-2,-1))\n",
        "    scale = torch.tensor(query.shape[-1],dtype=float)\n",
        "    logits = matmul / torch.sqrt(scale)\n",
        "    if mask is not None:\n",
        "        logits += (mask.float() * -1e9)\n",
        "    \n",
        "    attention_weights = F.softmax(logits,dim = -1)\n",
        "    output = torch.matmul(attention_weights,value)\n",
        "    return output"
      ],
      "metadata": {
        "id": "ZH0iIOkajtnZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Transformer Model\n",
        "\n",
        "This model has eight subparts:\n",
        "1. Embedder class\n",
        "2. Positional encoding class\n",
        "3. Multihead attention class\n",
        "4. Feed forward neural network \n",
        "5. Encoder classes (encoder layer + encoder + encoder block)\n",
        "6. Decoder classes (decoder layer + decoder)\n",
        "7. Convolutional neural network\n",
        "8. Transformer class"
      ],
      "metadata": {
        "id": "-mDVpjSpkXol"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note: Inspiration for certain core transformer implementation segments from \"Attention is All You Need\" paper are from: "
      ],
      "metadata": {
        "id": "iUz_ehW491G1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Embedder"
      ],
      "metadata": {
        "id": "cEWCN4AIkg4V"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Embedder(nn.Module):\n",
        "    '''Input embedding layer of size vocab_size * dimensionality\n",
        "    of word embedding'''\n",
        "    def __init__(self, vocab_size, d_model):\n",
        "        super().__init__()\n",
        "        self.embed = nn.Embedding(vocab_size, d_model)\n",
        "    \n",
        "    def forward(self,x):\n",
        "        return self.embed(x)"
      ],
      "metadata": {
        "id": "5aWiLZRVkmHj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Positional Encoding"
      ],
      "metadata": {
        "id": "SPbJdMWrkoew"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "    '''Transformers are not sequential so positional encoding\n",
        "    gives some sequentiality to sentence'''\n",
        "\n",
        "    def __init__(self, d_model, dropout=0.1, max_len=250):\n",
        "        super(PositionalEncoding, self).__init__()\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "        self.d_model = d_model\n",
        "\n",
        "        pe = torch.zeros(max_len, d_model)\n",
        "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, d_model, 2).float() \\\n",
        "                            * (-math.log(10000.0) / d_model))\n",
        "        pe[:, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 1::2] = torch.cos(position * div_term)\n",
        "        pe = pe.unsqueeze(0)\n",
        "        self.pe = pe\n",
        "\n",
        "    def forward(self, x):\n",
        "        x *= math.sqrt(self.d_model)\n",
        "        x +=  self.pe[:,:x.size(1)].cuda()\n",
        "        return self.dropout(x)"
      ],
      "metadata": {
        "id": "5T6tGB-MkstD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Multihead Attention"
      ],
      "metadata": {
        "id": "nHhkCMPKk3u0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "    '''Divides d_model into heads and\n",
        "    applies attention to each layer with helper \n",
        "    function scaled_dot_product_attention'''\n",
        "\n",
        "    def __init__(self, heads, d_model):\n",
        "        super().__init__()\n",
        "        self.heads = heads\n",
        "        self.d_model = d_model\n",
        "        print(d_model, self.heads)\n",
        "        assert d_model % self.heads == 0\n",
        "\n",
        "        self.d_h = self.d_model // self.heads\n",
        "\n",
        "        self.q_dense = nn.Linear(d_model,d_model)\n",
        "        self.k_dense = nn.Linear(d_model,d_model)\n",
        "        self.v_dense = nn.Linear(d_model,d_model)\n",
        "\n",
        "        self.out = nn.Linear(d_model,d_model)\n",
        "    \n",
        "    def forward(self, q, k, v, mask = None):\n",
        "        \n",
        "        # batch_size\n",
        "        bs = q.size(0)\n",
        "\n",
        "        k = self.k_dense(k).view(bs, -1, self.heads, self.d_h)\n",
        "        q = self.q_dense(q).view(bs, -1, self.heads, self.d_h)\n",
        "        v = self.v_dense(v).view(bs, -1, self.heads, self.d_h)\n",
        "\n",
        "        k = k.transpose(1,2)\n",
        "        q = q.transpose(1,2)\n",
        "        v = v.transpose(1,2)\n",
        "\n",
        "        scores = scaled_dot_product_attention(q,k,v,mask)\n",
        "        \n",
        "        # concat each heads\n",
        "        concat = scores.transpose(1,2).contiguous()\\\n",
        "            .view(bs,-1,self.d_model)\n",
        "        \n",
        "        out = self.out(concat)\n",
        "\n",
        "        return out"
      ],
      "metadata": {
        "id": "u9JwdRiik_Ba"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Feed Forward Neural Network"
      ],
      "metadata": {
        "id": "v-mPIBULlDLj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class FeedForward(nn.Module):\n",
        "    '''Feed Forward neural network'''\n",
        "    def __init__(self,d_model,d_ff = 1000,dropout = 0.1):\n",
        "        super().__init__()\n",
        "        self.linear_1 = nn.Linear(d_model,d_ff)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.linear2 = nn.Linear(d_ff, d_model)\n",
        "\n",
        "    def forward(self,x):\n",
        "        x = F.relu(self.linear_1(x))\n",
        "        x = self.dropout(x)\n",
        "        x = self.linear2(x)\n",
        "        return x"
      ],
      "metadata": {
        "id": "R-c8VewWlHUq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Encoder"
      ],
      "metadata": {
        "id": "oEFSKx5BlVRb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class EncoderLayer(nn.Module):\n",
        "    '''Encoder layer of transformer \n",
        "    embedding -> positional_encoding -> attention\n",
        "     -> Feed Forward with skip connection'''\n",
        "    def __init__(self, d_model, heads, dropout = 0.1):\n",
        "        super().__init__()\n",
        "        self.norm_1 = nn.LayerNorm(d_model)\n",
        "        self.norm_2 = nn.LayerNorm(d_model)\n",
        "        self.attn = MultiHeadAttention(heads, d_model)\n",
        "        self.ff = FeedForward(d_model)\n",
        "        self.dropout_1 = nn.Dropout(dropout)\n",
        "        self.dropout_2 = nn.Dropout(dropout)\n",
        "\n",
        "    \n",
        "    def forward(self,x):\n",
        "        x1 = self.norm_1(x)\n",
        "        x1 = x + self.dropout_1(self.attn(x1,x1,x1))\n",
        "        x2 = self.norm_2(x1)\n",
        "        x3 = x1 + self.dropout_2(self.ff(x2))\n",
        "        return x3\n",
        "\n",
        "class Encoder(nn.Module):\n",
        "    '''Cloning and making copies'''\n",
        "    def __init__(self, vocab_size, d_model, N, heads):\n",
        "        super().__init__()\n",
        "        self.N = N\n",
        "        self.embed = Embedder(vocab_size, d_model)\n",
        "        self.pe = PositionalEncoding(d_model)\n",
        "        self.layers = get_clones(EncoderLayer(d_model, heads), N)\n",
        "        self.norm = nn.LayerNorm(d_model)\n",
        "\n",
        "    def forward(self, src):\n",
        "        src[src < 0] = 0\n",
        "        src[src > 249] = 249\n",
        "\n",
        "        x = self.embed(src)\n",
        "        x = self.pe(x)\n",
        "\n",
        "        for i in range(self.N):\n",
        "            x = self.layers[i](x)\n",
        "        return self.norm(x) \n",
        "\n",
        "\n",
        "class EncoderBlock(nn.Module):\n",
        "      def __init__(self, vocab_size, d_model, num_layers, heads):\n",
        "        super().__init__()\n",
        "        self.encoder = Encoder(vocab_size,d_model,num_layers,heads)\n",
        "        self.conven = nn.Conv2d(12,1,1)\n",
        "\n",
        "      def forward(self, src):\n",
        "          src[src < 0] = 0\n",
        "          src[src > 249] = 249\n",
        "\n",
        "          src = torch.reshape(src,(1,250,12))\n",
        "          src = torch.transpose(src,0,2)\n",
        "\n",
        "          src = self.conven(src.float())\n",
        "          src = torch.squeeze(src,2)\n",
        "          \n",
        "          e_outputs = self.encoder(src.long())\n",
        "          return e_outputs # 1x250x120"
      ],
      "metadata": {
        "id": "E0BldElHlX7x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Decoder"
      ],
      "metadata": {
        "id": "2o-0zuSCliJp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class DecoderLayer(nn.Module):\n",
        "    '''Decoder layer - mha to notes embeddings -> add&norm -> mha to both notes embeddings and encoder output -> add&norm -> feed forward linear -> add&norm'''\n",
        "    def __init__(self, d_model, heads, dropout = 0.1):\n",
        "        super().__init__()\n",
        "        self.norm_1 = nn.LayerNorm(d_model)\n",
        "        self.norm_2 = nn.LayerNorm(d_model)\n",
        "        self.norm_3 = nn.LayerNorm(d_model)\n",
        "\n",
        "        self.dropout_1 = nn.Dropout(dropout)\n",
        "        self.dropout_2 = nn.Dropout(dropout)\n",
        "        self.dropout_3 = nn.Dropout(dropout)\n",
        "        \n",
        "        self.attn_1 = MultiHeadAttention(heads, d_model)\n",
        "        self.attn_2 = MultiHeadAttention(heads, d_model)\n",
        "        self.ff = FeedForward(d_model)\n",
        "\n",
        "    def forward(self, x, encoder_out):\n",
        "        x2 = self.norm_1(x)\n",
        "        x = x + self.dropout_1(self.attn_1(x2, x2, x2))\n",
        "        x2 = self.norm_2(x)\n",
        "        x = x + self.dropout_2(self.attn_2(x2, encoder_out, encoder_out))\n",
        "        x2 = self.norm_3(x)\n",
        "        x = x + self.dropout_3(self.ff(x2))\n",
        "        return x\n",
        "\n",
        "class Decoder(nn.Module):\n",
        "    '''Decoder module of transformer, running n sequential decoder segments - takes input of both encoder output and text embeddings'''\n",
        "    def __init__(self, d_model, vocab_size, N, heads):\n",
        "        super().__init__()\n",
        "        self.N = N\n",
        "        self.embed = Embedder(vocab_size, d_model)\n",
        "        self.pe = PositionalEncoding(d_model)\n",
        "        self.layers = get_clones(DecoderLayer(d_model, heads), N)\n",
        "        self.norm = nn.LayerNorm(d_model)\n",
        "\n",
        "    def forward(self, x, encoder_out):\n",
        "        for i in range(self.N):\n",
        "            x = self.layers[i](x, encoder_out)\n",
        "        return self.norm(x)"
      ],
      "metadata": {
        "id": "Y8g8kV6wljyB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Convolutional Neural Network"
      ],
      "metadata": {
        "id": "wJEyV68IlpGs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class CNN(nn.Module):\n",
        "    '''Convolutional output network for classification task - takes in summed input sequence and decoder output'''\n",
        "    def __init__(self):\n",
        "        super(CNN, self).__init__()\n",
        "        self.name = \"cnn\"\n",
        "        self.conv1 = nn.Conv2d(1, 5, 3)\n",
        "        self.pool = nn.MaxPool2d(2, 2)\n",
        "        self.conv2 = nn.Conv2d(5, 1, 3)\n",
        "        self.conv3 = nn.Conv2d(1, 1, 3)\n",
        "        self.fc1 = nn.Linear(377, 128)\n",
        "        self.fc2 = nn.Linear(128, 16)\n",
        "        self.fc3 = nn.Linear(16,5)\n",
        "    def forward(self, x):\n",
        "        x = self.pool(F.relu(self.conv1(x)))\n",
        "        x = self.pool(F.relu(self.conv2(x)))\n",
        "        x = self.pool(F.relu(self.conv3(x)))\n",
        "        x = x.view(x.size(0),-1)\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        x = x.squeeze(1) # Flatten to [batch_size]\n",
        "        return x"
      ],
      "metadata": {
        "id": "vfPKlcNal3w4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Transformer Model"
      ],
      "metadata": {
        "id": "tBCo8vtol-gT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Transformer(nn.Module):\n",
        "    '''Overall Transformer architecture'''\n",
        "    def __init__(self, vocab_size, d_model, num_layers, heads):\n",
        "        super().__init__()\n",
        "        self.name = \"Transformer\"\n",
        "        self.encoderBlock = EncoderBlock(vocab_size,d_model,num_layers,heads)\n",
        "        self.transform = torchvision.transforms.Resize((250,120))\n",
        "        self.decoder = Decoder(d_model,vocab_size,num_layers,heads)\n",
        "        self.conv_out = CNN()\n",
        "        self.fc = nn.Linear(30000, 5)\n",
        "        self.softmax = nn.Softmax(1)\n",
        "        self.transform_cnn = torchvision.transforms.Resize((250,120))\n",
        "        self.norm = nn.LayerNorm((1,250,120))\n",
        "\n",
        "    def forward(self, src, txt=None):\n",
        "        src = ((src + 1.68)*100).long() #scale\n",
        "        \n",
        "        e_outputs = self.encoderBlock(src.cuda()) #encoder\n",
        "        \n",
        "        # text embeddings\n",
        "        txt = torch.unsqueeze(txt, 0) \n",
        "        txt = self.transform(txt)\n",
        "        \n",
        "        d_output = self.decoder(txt, e_outputs.cuda()) #decoder\n",
        "        d_output = torch.squeeze(d_output,0)\n",
        "\n",
        "        src = self.transform_cnn(src.float())\n",
        "        return self.conv_out(self.norm(src + d_output)) #output"
      ],
      "metadata": {
        "id": "Uf5jitpRmGNX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training"
      ],
      "metadata": {
        "id": "SwrINy7RmrY3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Labels"
      ],
      "metadata": {
        "id": "iOCpzfVznJyu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#{['NORM']:0,['MI']:1,['STTC']:2,['CD']:3,['HYP']:4, ['HYP','MI']:5, ['HYP','CD']:6,\n",
        "#                  ['HYP','STTC']:7,['MI','CD']:8,['MI','STTC']:9,['CD','STTC']:10,['HYP','MI','STTC']:11,\n",
        "#                  ['HYP','MI','CD']:12, ['HYP','MI','CD','STTC']:13}\n",
        "W = np.array([[1,0,0,0,0],[0,1,0,0,0],[0,0,1,0,0],[0,0,0,1,0],[0,0,0,0,1],[0,1,0,0,1],[0,0,0,1,1],[0,0,1,0,1],[0,1,0,1,0],\n",
        "              [0,1,1,0,0],[0,0,1,1,0],[0,1,1,0,1],[0,1,0,1,1],[0,1,1,1,0],[0,0,1,1,1],[0,1,1,1,1],[1,0,0,1,0],[1,0,1,1,0],[1,0,1,0,0],\n",
        "              [1,0,0,0,1],[1,1,0,0,0],[1,1,0,0,1],[1,0,0,1,1],[1,1,0,1,1],[1,1,1,1,1],[1,0,1,0,1],[1,1,1,0,0],[1,1,1,0,1],[1,0,1,1,1],[1,1,1,1,0]])"
      ],
      "metadata": {
        "id": "yATOiT9vnIwE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Evaluation"
      ],
      "metadata": {
        "id": "BpT54YQRnfTk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(net, loader, criterion, batch_size):\n",
        "  accurate = 0\n",
        "  total = 0\n",
        "\n",
        "  total_err = 0\n",
        "  total_loss = 0\n",
        "  total_epoch = 0\n",
        "  enum = 1\n",
        "  for images, labels, note in iter(loader):\n",
        "    #############################################\n",
        "    #To Enable GPU Usage\n",
        "    if use_cuda and torch.cuda.is_available():\n",
        "      images = images.float().cuda()\n",
        "      labels = labels.cuda()\n",
        "      note = note.cuda()\n",
        "    #############################################\n",
        "    # images = torch.reshape(images,(1,1000,12))\n",
        "    # images = torch.transpose(images,0,2)\n",
        "    #outputs = net(torch.Tensor(images).to(torch.int64), src_mask, note)\n",
        "    output = net(images,note)\n",
        "\n",
        "    values, indices = torch.max(output,dim=1)\n",
        "    if use_cuda and torch.cuda.is_available():\n",
        "      loss = criterion(output,torch.tensor(W[labels]).float().cuda().unsqueeze(0))\n",
        "    else:\n",
        "      loss = criterion(output,torch.tensor(W[labels]).float().unsqueeze(0))\n",
        "\n",
        "    total_err += int(((W[labels][indices]) == 0))\n",
        "    total_loss += loss.item()\n",
        "    total_epoch += len(labels)\n",
        "    enum = enum + 1\n",
        "  err = float(total_err) / total_epoch\n",
        "  loss = float(total_loss) / (enum+1)\n",
        "  return err, loss"
      ],
      "metadata": {
        "id": "qij_Tn4Mnhf7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training"
      ],
      "metadata": {
        "id": "4fYYD5p_njNM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train_net(net,train_data,val_data,batch_size=1,learning_rate=5.0,num_epochs=3):\n",
        "  train_loader, val_loader = get_data_loader(train_data,val_data,batch_size)\n",
        "  net.train()  # turn on train mode\n",
        "  total_loss = 0.\n",
        "  #criterion = nn.CrossEntropyLoss()\n",
        "  criterion = nn.BCEWithLogitsLoss()\n",
        "  optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
        "  scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)\\\n",
        "\n",
        "  train_err = np.zeros(num_epochs)\n",
        "  train_loss = np.zeros(num_epochs)\n",
        "  val_err = np.zeros(num_epochs)\n",
        "  val_loss = np.zeros(num_epochs)\n",
        "  \n",
        "  # training\n",
        "  epoch = 0 # the number of iterations\n",
        "  for epoch in range(0, num_epochs):\n",
        "    total_train_loss = 0.0\n",
        "    total_train_err = 0.0\n",
        "    total_epoch = 0\n",
        "    enum = 0\n",
        "    for images, labels, note in iter(train_loader):\n",
        "      optimizer.zero_grad()\n",
        "      #############################################\n",
        "      #To Enable GPU Usage\n",
        "      if use_cuda and torch.cuda.is_available():\n",
        "        images = images.float().cuda()\n",
        "        labels = labels.cuda()\n",
        "        note = note.cuda()\n",
        "      #############################################\n",
        "      \n",
        "      output = model(images,note)\n",
        "      # print(output,labels)\n",
        "      values, indices = torch.max(output,dim=1)\n",
        "      \n",
        "      # print(W[labels.cpu()],indices.cpu(),output.cpu(), nn.functional.normalize(torch.sigmoid(output.cpu())))\n",
        "      if use_cuda and torch.cuda.is_available():\n",
        "        loss = criterion(output,torch.tensor(W[labels]).float().cuda().unsqueeze(0))\n",
        "      else:\n",
        "        loss = criterion(output,torch.tensor(W[labels]).float().unsqueeze(0))\n",
        "      # print(loss.cpu())\n",
        "     # if(W[labels.cpu()][indices.cpu()] == 1):\n",
        "      #  loss = criterion(output, indices)\n",
        "     # else:\n",
        "      #  loss = criterion(output, torch.tensor(np.argmax(W[labels])).cuda().unsqueeze(0))\n",
        "      \n",
        "      a = list(net.parameters())[0].clone()\n",
        "  \n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "\n",
        "      b = list(net.parameters())[0].clone()\n",
        "      # if (torch.equal(a.data, b.data)): \n",
        "      #   print(\"weights not updating\")\n",
        "\n",
        "      #print(list(model.parameters())[0].grad)\n",
        "      #for name, param in net.named_parameters():\n",
        "        #if param.requires_grad:\n",
        "          #print(name, param.data)\n",
        "\n",
        "      #0\n",
        "      #01111\n",
        "      #total_train_err = torch.sum((nn.functional.normalize(torch.sigmoid(output.cpu()))) - torch.tensor(W[labels]))\n",
        "      total_train_err += int(((W[labels][indices]) == 0))\n",
        "      total_train_loss += loss.item()\n",
        "      total_epoch += len(labels)\n",
        "      enum = enum + 1\n",
        "    train_err[epoch] = float(total_train_err) / total_epoch\n",
        "    train_loss[epoch] = float(total_train_loss) / (enum+1)\n",
        "    val_err[epoch], val_loss[epoch] = evaluate(net, val_loader, criterion, batch_size)\n",
        "    print((\"Epoch {}: Train Error: {}, Train loss: {} |\"+\n",
        "               \"Validation Error: {}, Validation loss: {}\").format(\n",
        "                   epoch + 1,\n",
        "                   train_err[epoch],\n",
        "                   train_loss[epoch],\n",
        "                   val_err[epoch],\n",
        "                   val_loss[epoch]))\n",
        "    if (torch.equal(a.data, b.data)): \n",
        "        print(\"weights not updating\")\n",
        "    model_path = \"model_{0}_lr{1}_epoch{2}\".format(model.name, learning_rate, epoch)\n",
        "    torch.save(model.state_dict(), model_path)\n",
        "\n",
        "  # plotting\n",
        "  plt.title(\"Training Curve\")\n",
        "  plt.plot(np.arange(1,num_epochs+1,1), train_loss, label=\"Train\")\n",
        "  plt.plot(np.arange(1,num_epochs+1,1), val_loss, label=\"Validation\")\n",
        "  plt.xlabel(\"Iterations\")\n",
        "  plt.ylabel(\"Loss\")\n",
        "  plt.legend(loc='best')\n",
        "  plt.show()\n",
        "\n",
        "  plt.title(\"Training Curve\")\n",
        "  plt.plot(np.arange(1,num_epochs+1,1), train_err, label=\"Train\")\n",
        "  plt.plot(np.arange(1,num_epochs+1,1), val_err, label=\"Validation\")\n",
        "  plt.xlabel(\"Iterations\")\n",
        "  plt.ylabel(\"Training Error\")\n",
        "  plt.legend(loc='best')\n",
        "  plt.show()"
      ],
      "metadata": {
        "id": "1WMgw4C2n0li"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Execution"
      ],
      "metadata": {
        "id": "e12gKxAQn2ac"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Instantiate the Model"
      ],
      "metadata": {
        "id": "7p6vlMOaoBjK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ntokens = 250  # size of vocabulary\n",
        "emsize = 120  # embedding dimension\n",
        "d_hid = 120  # dimension of the feedforward network model in nn.TransformerEncoder\n",
        "nlayers = 6  # number of nn.TransformerEncoderLayer in nn.TransformerEncoder\n",
        "nhead = 12  # number of heads in nn.MultiheadAttention\n",
        "dropout = 0.2  # dropout probability\n",
        "model = Transformer(ntokens, emsize, nlayers, nhead)"
      ],
      "metadata": {
        "id": "Crlbf1Avn77a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Run the Model"
      ],
      "metadata": {
        "id": "577k5MrvoEz6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "use_cuda = True\n",
        "if use_cuda and torch.cuda.is_available():\n",
        "  model.cuda()\n",
        "  print('CUDA is available!  Training on GPU ...')\n",
        "else:\n",
        "  print('CUDA is not available.  Training on CPU ...')\n",
        "\n",
        "train_net(model,train_data,val_data,batch_size=1,learning_rate=0.0001,num_epochs=20)"
      ],
      "metadata": {
        "id": "YYVa16FUoAHV"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}